import sys
import pickle
import argparse
import itertools
import collections
from pathlib import Path
from heapq import heapify, heappush, heappop
from typing import Union, Any, Dict, List, Tuple
from PyQt5.QtCore import QDir

import pandas
import pandas as pd
from colorama import Fore, Style
from PyQt5.QtWidgets import *

__author__ = "Jack Wang"
__date__ = "2021/12/19"
sys.path.append(Path(__file__).resolve().parent.joinpath("./ui").__str__())


class _EOFSymbol:
    """
    Notes:
        _EOFSymbol 是内置的文件结尾的类，用于在编码的时候标志文件的结尾。
        由于在进行哈夫曼树的构建的时候需要选出来权最小的节点，因此需要重写比较的魔术方法。
        默认文档结尾是频率最少的字符，因此和任何数字比较
    """

    def __repr__(self):
        return "_EOF"

    def __lt__(self, other):
        return True

    def __gt__(self, other):
        return False

    def __eq__(self, other):
        return self.__class__ == other.__class__

    def __hash__(self):
        return hash(self.__class__)


_EOF = _EOFSymbol()


def _guess_concat_function(data):
    """
    Notes:
        To support multiple types of input, e.g., str, bytes, list, there must be a
        function that return the needed concat function
    Returns:
        pass
    """
    return {type(u""): u"".join, type(b""): bytes}.get(type(data), list)


def ensure_dir(path: Union[str, Path]) -> Path:
    path: Path = path if isinstance(path, Path) else Path(path)
    if not path.exists():
        path.mkdir(parents=True)
    assert path.is_dir()
    return path


class _HuffmanCoder:
    """
    Notes:
        _HuffmanCoder, which provide encode, decode, serialize, deserialize, save_code, load_code functionalities.
    Args:

    """

    def __init__(self, code_table: Dict[Any, Tuple[int, int]], concat=list, check=True, eof=_EOF):
        """
        Notes:
            initialize method of _HuffmanCoder with given code table
        Args:
            code_table (Dict[Any, Tuple[int, int]]): code table of all symbols in the encoding sequence, create automatically by
                HuffmanCoderFactory, items are Tuple[Symbol, [depth, bit]]
            concat: concat method of symbols, will be automatically determined in HuffmanCoderFactory
            check: valid code table's format
            eof: eof symbol, leave it alone
        """
        self._table = code_table
        self._concat = concat
        self._eof = eof
        if check:
            assert isinstance(self._table, dict), f"{Fore.RED}Code table need to be a dict!{Style.RESET_ALL}"
            assert all(
                isinstance(b, int) and b >= 1 and isinstance(v, int) and v >= 0 for (b, v) in self._table.values()
            ), f"{Fore.RED}Code table internal format Error!{Style.RESET_ALL}"

    def get_code_table(self) -> Dict[Any, Tuple[int, int]]:
        """
        Notes:
            get_code_table returns code table of input sequence
        Returns:
            self._table
        """
        return self._table

    def print_code_table(self, verbose=False) -> pandas.DataFrame:
        """
        Notes:
            print_code_table is used to print code table
        Args:
            verbose (bool): if True, then print the table in a pretty format
        Returns:
            pd.DataFrame
        """
        pd.options.display.max_rows = None
        columns = list(zip(*list(zip(*itertools.chain(
            [("Symbol", "Huff Code", "BitSize")], (
                (repr(symbol), bin(val)[2:].rjust(bits, "0"), str(bits))
                for symbol, (bits, val) in self._table.items()
            )
        )))))
        df = pd.DataFrame(columns[1:], columns=columns[0])
        if verbose:
            print(df)
        return df

    def _encode_streaming(self, raw_sequence):
        size = 0
        buffer = 0
        for s in raw_sequence:
            bits, values = self._table[s]
            buffer = (buffer << bits) + values
            size += bits
            while size >= 8:
                byte = buffer >> (size - 8)
                yield byte
                buffer = buffer - (byte << (size - 8))
                size -= 8
        if size > 0:
            bit, value = self._table[self._eof]
            buffer = (buffer << bit) + value
            size += bit
            if size > 8:
                byte = buffer >> (size - 8)
            else:
                byte = buffer << (8 - size)
            yield byte

    def _decode_streaming(self, encoded_sequence):
        # Reverse lookup table: map (bitsize, value) to symbols
        lookup = {(b, v): s for s, (b, v) in self._table.items()}

        buffer = 0
        size = 0
        for byte in encoded_sequence:
            for m in [128, 64, 32, 16, 8, 4, 2, 1]:
                buffer = (buffer << 1) + bool(byte & m)
                size += 1
                if (size, buffer) in lookup:
                    symbol = lookup[size, buffer]
                    if symbol == self._eof:
                        return
                    yield symbol
                    buffer = 0
                    size = 0

    def encode(self, raw_seq):
        """
        Notes:
            encode is used to encode the sequence that build the huffman tree
        Args:
            raw_seq: sequence that build the huffman tree
        Returns:
            bytes
        """
        return bytes(self._encode_streaming(raw_seq))

    def decode(self, encoded_seq, concat=None):
        return (concat or self._concat)(self._decode_streaming(encoded_seq))

    def save(self, path: Union[str, Path], metadata: Any = None):
        code_table = self.get_code_table()
        data = {
            "code_table": code_table,
            "type": type(self),
            "concat": self._concat,
        }
        if metadata:
            data['metadata'] = metadata
        path = Path(path)
        ensure_dir(path.parent)
        with path.open(mode='wb') as f:
            pickle.dump(data, file=f)

    @staticmethod
    def load(path: Union[str, Path]) -> '_HuffmanCoder':
        path = Path(path)
        with path.open(mode='rb') as f:
            data = pickle.load(f)
        cls = data['type']
        assert issubclass(cls, _HuffmanCoder)
        code_table = data['code_table']
        return cls(code_table, concat=data['concat'])


class HuffmanCoderFactory(_HuffmanCoder):
    """
    Notes:
        HuffmanCodecFactory is the class that create HuffmanCoder from different type of inputs
    Methods:
        from_frequencies: create HuffmanCoder from frequency table
    """

    @classmethod
    def from_frequencies(cls, frequencies: Dict[Any, int], concat=None, eof=_EOF):
        """
        Notes:
            from_frequencies creates huffman codec by symbol-frequency table/mapping.
        Args:
            frequencies (Dict[Any, int]): Symbols and its frequency, symbols can be str, bytes or int, etc.
            concat (Union[None]): concat method of symbols, will be determined by the function if argument is not
                specified
            eof (_EOFSymbol): leave it alone.
        Returns:
            __HuffmanCoder
        Examples:
            >>> huf_coder = HuffmanCoderFactory.from_frequencies({"a":29, "b":10, "c": 5})
            >>> type(huf_coder)
        """
        concat_function = concat if concat is not None else _guess_concat_function(next(iter(frequencies)))

        # build huffman tree node heap
        # each item: (frequency, [(symbol, (bitsize, value))], value equals which layer of the tree
        heap: List[Tuple[int, List[Tuple[Any, Tuple[int, int]]]]] = [(freq, [(symbol, (0, 0))]) for symbol, freq in
                                                                     frequencies.items()]
        # add eof
        if eof not in frequencies:
            heap.append((1, [(eof, (0, 0))]))

        heapify(heap)
        while len(heap) > 1:
            # get first 2 min as left and right child tree
            a: Tuple[int, List[Tuple[Any, Tuple[int, int]]]] = heappop(heap)
            b: Tuple[int, List[Tuple[Any, Tuple[int, int]]]] = heappop(heap)
            # merge child to form parent
            # parent frequencies adds together, left child add 0 ahead (do nothing) of previous bits
            # right add 1 ahead of previous bits
            merged = (
                a[0] + b[0],
                [(s, (n + 1, v)) for (s, (n, v)) in a[1]] + [(s, (n + 1, (1 << n) + v)) for (s, (n, v)) in b[1]]
            )
            heappush(heap, merged)

        # code table is root
        table = dict(heappop(heap)[1])
        return cls(table, concat=concat, check=False, eof=eof)

    @classmethod
    def from_sequence(cls, sequence: Union[List[Any], Tuple[Any], str, bytes]):
        """
        Notes:
            from_sequence will build a huffman tree from a sequence of symbol
        Args:
            sequence (Union[List, Tuple][Any]):
        Returns:
            __HuffmanCoder
        Examples:
            >>> seq = "a"*100 + "b"*29 + "c"*32
            >>> coder = HuffmanCoderFactory.from_sequence(seq)
        """
        frequencies = collections.Counter(sequence)
        return cls.from_frequencies(frequencies, concat=_guess_concat_function(sequence))


def encode(path: Union[str, Path], verbose: bool=True):
    assert (path := Path(path)).exists(), f"{Fore.RED}{path} not exists!{Style.RESET_ALL}"
    with path.open(mode="rb") as f:
        content = f.read()
    coder = HuffmanCoderFactory.from_sequence(content)
    byte = coder.encode(content)
    if not (out_fd := path.resolve().parent / (path.stem + "_huf")).exists():
        out_fd.mkdir(parents=True)
    with (main_path := out_fd.joinpath(path.stem + "_" + path.suffix[1:] + ".huf")).open("wb") as f:
        f.write(byte)
    coder.save(coder_path := out_fd.joinpath(path.stem + "_coder.huf"))

    b1 = path.stat().st_size
    b2 = main_path.stat().st_size
    b3 = coder_path.stat().st_size
    summary1, summary2, summary3, summary4 = [""]*4
    if verbose:
        summary1 = f"Origin file: {path.stem}{path.suffix}, byte size: {b1}"
        summary2 = f"Compressed file: {main_path.stem}{main_path.suffix}, byte size: {b2}"
        summary3 = f"HuffCoder file: {coder_path.stem}{coder_path.suffix}, byte size: {b3}"
        summary4 = f"{Fore.GREEN}Compression rate {b2}/{b1}={Style.BRIGHT}{round(b2/b1, 4)*100}%{Style.NORMAL}" +\
                f", saved {Style.BRIGHT}{b1-b2}{Style.RESET_ALL} bytes"
        print(s := "Summary".center(max([len(summary1), len(summary2), len(summary3), len(summary4)]), "="))
        print("\n".join((summary1, summary2, summary3, summary4)))
        print("="*len(s))
    return b1, b2, b3, coder


def decode(huf_path: Union[Path, str], coder_path: Union[Path, str], verbose: bool=True):
    huf_path, coder_path = Path(huf_path), Path(coder_path)
    coder = _HuffmanCoder.load(coder_path)
    with huf_path.open(mode="rb") as f:
        content = f.read()
    decoded = coder.decode(content)
    with (dec_path := huf_path.parent.joinpath((s := huf_path.stem.split("_"))[0] + f"_decode.{s[1]}")).open(mode="wb") as f:
        f.write(decoded)
    
    b1 = huf_path.stat().st_size
    b2 = coder_path.stat().st_size
    b3 = dec_path.stat().st_size
    if verbose:
        summary1 = f"Origin file: {huf_path.stem}{huf_path.suffix}, byte size: {b1}"
        summary2 = f"Coder file: {coder_path.stem}{coder_path.suffix}, byte size: {b2}"
        summary3 = f"Decompressed file: {dec_path.stem}{dec_path.suffix}, byte size: {b3}"
        summary4 = f"{Fore.GREEN}Decompression rate {b3}/{b1}={Style.BRIGHT}{round(b3/b1, 4)*100}%{Style.NORMAL}" +\
                f", lost {Style.BRIGHT}{b3-b1}{Style.RESET_ALL} bytes"
        print(s := "Summary".center(max([len(summary1), len(summary2), len(summary3), len(summary4)]), "="))
        print("\n".join((summary1, summary2, summary3, summary4)))
        print("="*len(s))
    return b1, b2, b3, coder



def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="哈夫曼压缩算法Python实现，可以对任意文件进行压缩，提供命令行以及GUI用户界面。作者：Jack Wang"
    )
    parser.add_argument("-s", "--src", dest="src_path", type=str, default=None, help="需要压缩的文件路径")
    parser.add_argument("-t", "--target", dest="target_path", type=str, default=None, help="需要解压文件所在的文件夹的路径")
    parser.add_argument("-v", "--verbose", dest="v", action="store_true", help="是否显示压缩的摘要信息")
    parser.add_argument("-g", "--gui", dest="gui", action="store_true", help="是否以GUI方式显示")
    return parser.parse_args()


class QtGUIHuff(QMainWindow):
    def __init__(self) -> None:
        super().__init__()
        self.initUI()

    def initUI(self):
        from Ui_main import Ui_Form
        self.ui = Ui_Form()
        self.ui.setupUi(self)
        self.setWindowTitle("哈夫曼压缩/解压缩程序")
        self.ui.compress.clicked.connect(self.compress)
        self.ui.decompress.clicked.connect(self.decompress)

    def fileDialog(self):
        dialog = QFileDialog()
        dialog.setFileMode(QFileDialog.AnyFile)
        dialog.setFilter(QDir.Files)
        if dialog.exec():
            filename = dialog.selectedFiles()
            return Path(filename[0])
        else:
            return False

    def compress(self):
        f_path = self.fileDialog()
        self.ui.textEdit.setReadOnly(True)
        s = ["开始压缩"]
        self.ui.textEdit.setText("\n".join(s))
        result = encode(f_path)
        summary1 = f"Origin file: {f_path.stem}{f_path.suffix}, byte size: {result[0]}"
        summary2 = f"Compressed file:  byte size: {result[1]}"
        summary3 = f"HuffCoder file: , byte size: {result[2]}"
        summary4 = f"Compression rate {result[1]}/{result[0]}={round(result[1]/result[0], 4)*100}%" + \
                f", saved {result[0] - result[1]} bytes"
        self.ui.textEdit.setText("\n".join([*s, summary1, summary2, summary3, summary4, "压缩结束"]))
        self.ui.textEdit_2.setText(result[-1].print_code_table(verbose=False).__str__())

    def decompress(self):
        f_path = self.fileDialog()
        if f_path.suffix != ".huf":
            QMessageBox.warning(self, "警告", "请选择正确的压缩文件路径", QMessageBox.Yes, QMessageBox.Yes)
            return False
        
        if len(ps := list(Path(f_path).parent.glob("*.huf"))) == 2:
            print(ps)
            for p_idx in range(len(ps)):
                if "coder" in ps[p_idx].stem:
                    cp_idx = p_idx
        else:
            QMessageBox.warning(self, "警告", "缺少必要的解压缩文件", QMessageBox.Yes, QMessageBox.Yes)
            return False
        s = ["开始解压"]
        import time
        time.sleep(1)
        self.ui.textEdit.setText("\n".join(s))
        result = decode(huf_path:=ps[1-cp_idx], coder_path:=ps[cp_idx])
        summary1 = f"Origin file: {huf_path.stem}{huf_path.suffix}, byte size: {result[0]}"
        summary2 = f"Coder file: {coder_path.stem}{coder_path.suffix}, byte size: {result[1]}"
        summary3 = f"Decompressed file: byte size: {result[2]}"
        summary4 = f"Decompression rate {result[2]}/{result[0]}={round(result[2]/result[0], 4)*100}%" +\
                f", lost {result[2] - result[0]} bytes"
        self.ui.textEdit.setText("\n".join([*s, summary1, summary2, summary3, summary4, "解压结束"]))
        self.ui.textEdit_2.setText(result[-1].print_code_table(verbose=False).__str__())



def main(args: argparse.Namespace):
    if not args.gui:
        assert 0 < int(args.src_path is not None) + int(args.target_path is not None), f"{Fore.RED}请选择压缩或者解压，使用-h查询用法{Style.RESET_ALL}"
    if not args.gui:
        print(f"{Fore.GREEN}命令行工作模式{Style.RESET_ALL}")
        if args.src_path is not None:
            print(f"开始压缩 {args.src_path}")
            if args.v:
                print(encode(Path(args.src_path))[-1].print_code_table(verbose=False))
            else:
                encode(Path(args.src_path), verbose=False)
            print("压缩完成")
        elif args.target_path is not None:
            if len(ps := list(Path(args.target_path).glob("*.huf"))) == 2:
                print(ps)
                for p_idx in range(len(ps)):
                    if "coder" in ps[p_idx].stem:
                        cp_idx = p_idx
            else:
                assert False, f"{Fore.RED}无效的路径，没有找到压缩文件{Style.RESET_ALL}"
            print(f"开始解压 {args.target_path}")
            if args.v:
                print(decode(huf_path=ps[1-cp_idx], coder_path=ps[cp_idx])[-1].print_code_table(verbose=False))
            else:
                decode(huf_path=ps[1-cp_idx], coder_path=ps[cp_idx], verbose=False)
            print("解压完成")
    else:
        app = QApplication(sys.argv)
        m = QtGUIHuff()
        m.show()
        sys.exit(app.exec_())




if __name__ == "__main__":
    main(parse_args())
